#!/usr/bin/python

import csv

from tqdm import tqdm

from util.dct import safe_save
from util.log import logger


class CSVFile:
    def __init__(self, file, encoding):
        self.path = file
        self.file = None
        self.encoding = encoding

    def open(self, mode):
        if self.file is not None:
            self.file.close()
        if mode == 'r':
            self.file = open(self.path, 'r', encoding=self.encoding)
        elif mode == 'w':
            self.file = open(self.path, 'w', encoding=self.encoding, newline='')

    def read_dict(self, auto_close=True):
        if self.file is None:
            self.open('r')
        reader = csv.reader(self.file)
        dct = {}
        logger.info("开始读取cvs文件中的数据：{}".format(self.path))
        for item in tqdm(reader):
            if len(item) > 1:
                safe_save(dct, item[0].encode('utf-8').decode('utf-8-sig'), item[1].encode('utf-8').decode('utf-8-sig'))
        if auto_close:
            self.close()
        logger.info("读取数据完毕")
        return dct

    def write_dict(self, dct, auto_close=True):
        if self.file is None:
            self.open('w')
        writer = csv.writer(self.file)
        logger.info("开始写入cvs文件数据：{}".format(self.path))
        for key in tqdm(dct):
            writer.writerow([key, dct[key]])
        if auto_close:
            self.close()
        logger.info("写入数据完毕")

    def close(self):
        if self.file is not None:
            self.file.close()
